/*
 * $AIST_Release: 0.9.0 $
 * Copyright 2011 Information Technology Research Institute, National
 * Institute of Advanced Industrial Science and Technology
 * 
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *    http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package parser.analyzer

import scala.collection.JavaConversions._
import scala.collection.mutable.{Map => MMap}
import scala.collection.mutable.{Set => MSet}
import scala.collection.mutable.ListBuffer
import scala.Stream

import com.google.protobuf.DescriptorProtos._
import com.google.protobuf.Descriptors._

import aggr.{SerializedTypes => ST}

import parser.TreeWalker
import parser.DataTypes._
import parser.Nodes._
import parser.Results._

/**
 * Makes information to serialize the table contents.
 */
object TableSerializationInfoMaker {

  val dummyPackageName = "sawzallclone"
  val dummyFileName    = "dummy82qhfuabuhg982qhg.proto"

  /**
   * Results of generate().
   */
  case class State(ast: List[Stmt], fd: FileDescriptor, anonymousMap: Map[DataType, Descriptor])

  type IType = Either[ST.SerializedType, DescriptorProto]

  def generate(ast: List[Stmt], typeMap: Map[Int, DataType]): Result[State] = {

    makeSerializationInfo(ast, typeMap).flatMap { 
      case (serializedTypeMap, anonymousFileDescriptor, anonymousMap) =>
        burySerializationInfo(ast, serializedTypeMap).map { ast =>
          State(ast, anonymousFileDescriptor, anonymousMap)
        }
    }
  }

  private def makeSerializationInfo(ast: List[Stmt], typeMap: Map[Int, DataType]) = {
    var sinfo    = ListBuffer[(IType, List[IType], IType)]()
    val msgMaker = new MessageMaker(typeMap)

    val r = TreeWalker.walkTree(classOf[Node], ast) {
      case s@DefTableStmt(_, TableInfo(_, arg, keyType, valType, opts)) =>
        val keyR = msgMaker.make(mkKeyTupleType(keyType.map{_.dataType}))
        val valR = msgMaker.make(mkRealValueType(valType.dataType, opts))

        // Arguments type is only int, now
        val argsType = arg.toList.map { _ => Left(ST.Int64Type) }

        (keyR ~ valR).map { case k ~ v =>
          sinfo.append((k, argsType, v))
        }
    }

    def iType2SerializedType(itype: IType): ST.SerializedType = itype match {
      case Left(st)     => st
      case Right(proto) => 
        ST.MessageType(msgMaker.anonymousFileDescriptor.findMessageTypeByName(proto.getName))
    }

    r.mapConst {
      (
        List() ++ sinfo.map { case (k, as, v) => 
          TableSerializationInfo(iType2SerializedType(k), as.map(iType2SerializedType), iType2SerializedType(v)) 
        }, 
        msgMaker.anonymousFileDescriptor,
        msgMaker.anonymousMap
      )
    }
  }

  private def burySerializationInfo(ast: List[Stmt], sinfo: List[TableSerializationInfo]) = {
    var it = sinfo
    TreeWalker.mapTree(classOf[Node], ast) { case s@DefTableStmt(vars, tInfo) =>
        try     Success(DefTableStmtWithSerializationInfo(vars, tInfo, it.head).setPos(s.pos))
        finally it = it.tail
    }.map { _.asInstanceOf[List[Stmt]] }
  }


  def mkKeyTupleType(keys: List[DataType]): DataType = keys match {
    case List()  => EmptyType
    case List(x) => x
    case x       => PureTupleType(Stream.from(0).map{"member" + _ }.zip(keys).toList)
  }

  def mkRealValueTypeFromTuple(value: DataType, opts: List[(String, DataType)]): DataType = {
    opts match {
      case Nil => value
      case _   => PureTupleType( ("data__", value) :: opts)
    }
  }

  def mkRealValueType(value: DataType, opts: List[TableOption]): DataType = {
    mkRealValueTypeFromTuple(value, opts.map { x => (x.id, x.dataType)})
  }

  def mkArrayTuple(dataType: DataType) = {
    PureTupleType(List(("elems", dataType)))
  }
}

import TableSerializationInfoMaker._

/**
 * Anonymous PB message maker.
 */
private class MessageMaker(typeMap: Map[Int, DataType]) {

  private val prefix = "." + dummyPackageName + "."
  private var nameSeed = 1
  private val files = MSet[FileDescriptor]()
  private var protoTable = MMap[DataType, DescriptorProto]()
  
  /**
   * File descriptor fo anonymous messages.
   */
  lazy val anonymousFileDescriptor: FileDescriptor = {
    val proto = FileDescriptorProto.newBuilder()
        .addAllDependency(files.map{_.getName}) 
        .addAllMessageType(asCollection(protoTable.values))
        .setName(dummyFileName)
        .setPackage(dummyPackageName).build()
    FileDescriptor.buildFrom(proto, files.toArray)
  }

  /**
   * Map dataType to Descriptor.
   * This contains anonymous messages.
   */
  lazy val anonymousMap: Map[DataType, Descriptor] = {
    Map() ++ protoTable.mapValues { proto =>
      anonymousFileDescriptor.findMessageTypeByName(proto.getName)
    }
  }

  /**
   * Makes DescriptorProto table of type defined in Sawzall Clone source file(not proto files).
   */
  def make(dataType: DataType): Result[IType] = {

    def label(dataType: DataType) = {
      import FieldDescriptorProto.Label._
      dataType match {
        case ArrayType(_) => LABEL_REPEATED
        case _            => LABEL_REQUIRED
      }
    }

    def fromMember(name: String, dataType: DataType, nth: Int): Result[FieldDescriptorProto.Builder] = {
      val fld = FieldDescriptorProto.newBuilder()
      val nonArrayType = dataType match {
        case ArrayType(t) => t
        case t            => t
      }

      fromDataType(nonArrayType).flatMap { x => 
        fld.setNumber(nth)
        fld.setName(name)
        fld.setLabel(label(dataType))

        x match {
          case Left(st) => 
            fld.setType(st.protoType)
            st match {
              case ST.MessageType(desc) => fld.setTypeName("." + desc.getFullName)
              case ST.EnumType(desc)    => fld.setTypeName("." + desc.getFullName)
              case _ =>
            }
          case Right(st) =>
            fld.setType(FieldDescriptorProto.Type.TYPE_MESSAGE)
            fld.setTypeName(prefix + st.getName)
        }

        Success(fld)
      }
    }

    def fromTuple(t: TupleType): Result[DescriptorProto] = {
      protoTable.get(t) match {
        case Some(proto) => Success(proto)
        case None =>
          val builder = DescriptorProto.newBuilder()

          val mems = t.members.zipWithIndex.map { case ((name, dataType), n) => 
            fromMember(name, dataType, n+1).map { fld => 
              builder.addField(fld)
            }
          }

          mergeList(mems).mapConst {
            builder.setName("Anonymous%d".format(nameSeed))
            val proto = builder.build()

            // Update
            protoTable(t) = proto
            nameSeed += 1

            proto
          }
      }
    }

    def fromDataType(t: DataType): Result[IType] = t match {
      case TypeReference(id) => typeMap.get(id) match {
        case Some(dataType) => fromDataType(dataType)
        case None => error("Invalid type id %d".format(id))
      }
      case e: EnumType      => Success(Left(ST.EnumType(e.desc)))
      case m: MessageType   => Success(Left(ST.MessageType(m.desc)))
      case array: ArrayType => fromTuple(mkArrayTuple(array)).map { Right(_) }
      case tuple: TupleType => fromTuple(tuple).map { Right(_) }
      case _: MapType => failure(t.pos, "Map type is not emitted.")
      case IntType          => Success(Left(ST.Int64Type))
      case FloatType        => Success(Left(ST.DoubleType))
      case TimeType         => Success(Left(ST.Fixed64Type))
      case FingerPrintType  => Success(Left(ST.Fixed64Type))
      case StringType       => Success(Left(ST.StringType))
      case BytesType        => Success(Left(ST.BytesType))
      case EmptyType        => Success(Left(ST.EmptyType))
      case x: ImperfectType => error("Invalid type: " + x)
      case x: FunctionList  => error("Invalid type: " + x)
      case x: FunctionType  => error("Invalid type: " + x)
      case x: RefType       => error("Invalid type: " + x)
      case x: TableType     => error("Invalid type: " + x)
      case x: WildCardType  => error("Invalid type: " + x)
      case UnkownType       => error("Invalid type: " + t)
    }

    fromDataType(dataType)
  }
}
